from sys import stdout

print_intermediates = True
print_start_values  = True
print_init_matrices = True
print_check_matrix  = True

spu_madd    = lambda v1,v2,v3: [(i*j) + k for i,j,k in zip(v1,v2,v3)]

gen1        = lambda size    : [[i + j for j in range(size)] for i in range(size)]

gen2        = lambda size    : [[2*j - i + size for j in range(size)] for i in range(size)]

gen3        = lambda size    : [[0 for j in range(size)] for i in range(size)]

rotl        = lambda v, am   : [v[(i+am) % len(v)] for i in range(len(v))]

rot_all_l   = lambda m       : [rotl(v,1) for v in m]

shift_all_u = lambda m       : [m[(i+1) % len(m)] for i in range(len(m))]

shift_init  = lambda m, dim  : [[m[(i + j) % dim][i] for i in range(dim)] for j in range(dim)]

vec_to_s    = lambda v       : ' '.join(str(elem) for elem in v)

mat_to_s    = lambda m       : '\n'.join(vec_to_s(v) for v in m)

rot_init    = lambda m       : [rotl(m[i],i) for i in range(len(m))]		

printAll    = lambda a,b,c   : stdout.write("A matrix\n\n%s\n\nB matrix\n\n%s\n\nC matrix\n\n%s\n\n" %(mat_to_s(a), mat_to_s(b), mat_to_s(c)))

#does the normal non-vectorized matrix multiplication
def matmul(m1,m2,m3):
	for i in range(len(m1)):
		for j in range(len(m1)):
			for k in range(len(m1)):
				m3[i][k] += m1[i][j] * m2[j][k]
	return m3

#does the 'vectorized multiplication'
def matmul_vector(A,B,C):
	for i in range(4):
		for j in range(4):
			C[j] = spu_madd(A[j], B[j], C[j])
		A = rot_all_l(A)
		B = shift_all_u(B)
		
		if print_intermediates: 
			print
			print "---------- Matrices after %d rounds of multiplication ------------" % i
			print
			printAll(A,B,C)
			print
			print "------------------------------------------------------------------"
			print
			
	return C

#checks two matrices for equality
def eq(A,B):		
	for i in range(4):
		for j in range(4):
			if A[i][j] != B[i][j]:
				print "Not equal at %d,%d" % (i,j)
	
def main():
	######  NOTE ######
	#  using matrixs of size 4 to
	# simulate using vectors on the cell
	# spus

	#generate the A matrix to be multiplied
	matA = gen1(4)
	#generate the B matrix to be multiplied
	matB = gen2(4)
	
	#generates a C matrix initialized to all zeros
	matC = gen3(4)
	
	
	
	if print_start_values: 
		print
		print "---------- Starting Matrix Values Before Initialization ----------"
		print
		printAll(matA, matB, matC)
		print
		print "------------------------------------------------------------------"
		print
	
	A = rot_init(matA)
	B = shift_init(matB,4)
	C = matC
	
	if print_init_matrices: 
		print
		print "---------- Starting Matrix Values After Initialization -----------"
		print
		printAll(A, B, C)
		print
		print "------------------------------------------------------------------"
		print

	C = matmul_vector(A,B,C)

	check = matmul(gen1(4), gen2(4), gen3(4))
	
	eq(C,check)

	if print_check_matrix: 
		print
		print "---------- Matrix created with normal multiplication ----------"
		print
		print mat_to_s(check) 
		print
		print"----------------------------------------------------------------"
	
if __name__ == '__main__':
	main()
		
		
		
	